import cv2 as cv
import numpy as np
import socket
import struct

# pc is server
HOST = '219.228.135.234'
PORT = 7890
BUFFER_SIZE = 65535
a = True

def image_enhance(img, gamma):
    img = np.power(img / 255, gamma)
    return img

def image_enhance(img):
    (b, g, r) = cv.split(img)
    bH = cv.equalizeHist(b)
    gH = cv.equalizeHist(g)
    rH = cv.equalizeHist(r)
    out_img = cv.merge((bH, gH, rH))
    return out_img

def main():
    server=socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server.bind((HOST, PORT))
    print('wait for frames..')
    while True:
        data, address = server.recvfrom(BUFFER_SIZE)
        if len(data) == 1 and data[0] == 1:
            server.close()
            cv.destroyAllWindows()
            exit()
        if len(data) != 4:
            length = 0
        else :
            length = struct.unpack('i', data)[0]
        data, address = server.recvfrom(BUFFER_SIZE)
        if length != len(data):
            continue
        data = np.array(bytearray(data))
        img_decode = cv.imdecode(data, 1)
        img_decode = image_enhance(img_decode)
        cv.imshow('streamed frames', img_decode)
        if cv.waitKey(1) == 27:
            break

    server.close()
    cv.destroyAllWindows()

if __name__ == "__main__":
    main()